import os
import unittest

import torch
from config import Config
from torch import hub
from torch.backends import cudnn

from training.losses import TripletLoss

CONFIG = Config()
hub.set_dir(CONFIG['TORCH_HOME'])
os.environ["CUDA_VISIBLE_DEVICES"] = CONFIG['CUDA_VISIBLE_DEVICES']
torch.backends.cudnn.benchmark = True


class CompactNetTestCase(unittest.TestCase):

    def test_triplet_loss(self):
        y_pred = torch.ones(4, 1000)
        y_pred[2:] = -1.

        y = torch.ones(4)
        y[2:] = 0

        criterion = TripletLoss()
        loss = criterion(y_pred, y)
        self.assertEqual(0.0, loss.item())


if __name__ == '__main__':
    unittest.main()
